package HandshakeJoin;

import java.util.List;
import java.util.LinkedList;
import java.util.Collection;
import java.util.Calendar;
import java.util.concurrent.TimeUnit;
import java.math.RoundingMode;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimap;
import com.google.common.collect.LinkedListMultimap;
import com.google.common.base.Preconditions;
import org.apache.commons.collections.buffer.BoundedFifoBuffer;
import com.google.common.base.Joiner;
import org.apache.s4.core.ProcessingElement;
import org.apache.s4.core.Stream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class JoinWorker extends ProcessingElement
{
    transient private static final Logger logger = LoggerFactory
            .getLogger(JoinWorker.class);

    private static enum BSTType {
        RED_BLACK, SPLAY, UNKNOWN
    }

    private Stream<Message> downStreamRight = null;
    private Stream<Message> downStreamLeft = null;
    private Stream<JoinResult> downStreamResult = null;
    private BoundedFifoBuffer winR;
    private BoundedFifoBuffer winS;
    private LinkedList<Tuple> fwdR;
    private LinkedList<Tuple> fwdS;
    private Multimap<Integer, Tuple> mmapR;
    private Multimap<Integer, Tuple> mmapS;
    private BST<Integer, Tuple> bstR;
    private BST<Integer, Tuple> bstS;
    private int capWinR = 256;
    private int capWinS = 256;
    private int sizeNextWinR = 0;
    private int sizeNextWinS = 0;
    private boolean rightMost = false;
    private boolean leftMost = false;
    private long currTime = -1;
    private LoadShedder loadShedder = null;
    private boolean blockedTrans = false;
    private int blkSizeR = 16;
    private int blkSizeS = 16;
    private boolean multimapJoin = false;
    private boolean bstJoin = false;
    private BSTType bstType = BSTType.RED_BLACK;
    private long cmpCount;

    @Override
    protected void onCreate() {
        winR = new BoundedFifoBuffer(capWinR);
        winS = new BoundedFifoBuffer(capWinS);
        fwdS = Lists.newLinkedList();

        if (rightMost) {
            fwdR = Lists.newLinkedList();

            int thrhStartShedding = (int) (capWinS * 0.75);
            loadShedder = new LoadShedder(capWinS, thrhStartShedding, 0.3);
        }

        if (leftMost) {
            int thrhStartShedding = (int) (capWinR * 0.75);
            loadShedder = new LoadShedder(capWinR, thrhStartShedding, 0.3);
        }

        if (multimapJoin) {
            mmapR = LinkedListMultimap.create(capWinR);
            mmapS = LinkedListMultimap.create(capWinS);
        }

        if (bstJoin) {
            switch (bstType) {
            case SPLAY:
                bstR = new SplayBST<Integer, Tuple>();
                bstS = new SplayBST<Integer, Tuple>();
                break;

            case RED_BLACK:
                bstR = new RedBlackBST<Integer, Tuple>();
                bstS = new RedBlackBST<Integer, Tuple>();
                break;

            default:
                logger.error("Unknown BST type");
                break;
            }
        }

        cmpCount = 0;
    }

    public void onTime() {
        if (getSizeWindowR() > 0 || getSizeWindowS() > 0) {
            System.out.println(getStat());            
        }
    }

    public void onEvent(Message msg) {
        if (rightMost || leftMost) {
            currTime = Calendar.getInstance().getTimeInMillis();
        }

        if (msg.from() == Message.Direction.FROM_LEFT) {
            processLeft(msg);
        }

        if (msg.from() == Message.Direction.FROM_RIGHT) {
            processRight(msg);
        }

        if (blockedTrans) {
            forwardTupleBlocks();
        }
        else {
            forwardTuples();
        }
    }

    @Override
    protected void onRemove() {}

    private void processLeft(final Message msg) {
        if (msg.getType() == Message.Type.TUPLE) {
            if (leftMost) {
                if (loadShedder.shedIt(getSizeWindowR())) {
                    logger.info("Tuple R is shed");
                    return;
                }
            }

            Tuple r = msg.getTuple();

            joinWindowS(r);

            addToWindowR(r);

            ackLeft();
        }
        else if (msg.getType() == Message.Type.TUPLE_BLK) {
            TupleBlock rBlk = msg.getTupleBlock();

            for (Tuple r : rBlk) {
                joinWindowS(r);

                addToWindowR(r);
            }

            ackBlockLeft();
        }
        else if (msg.getType() == Message.Type.ACK) {
            sizeNextWinS = msg.getWindowSize();

            if (leftMost) {
                evictExpiredTuplesInWindowS();
            }
            else {
                removeLastInWindowS();
            }

            infoSizeChangedRight();
        }
        else if (msg.getType() == Message.Type.ACK_BLK) {
            sizeNextWinS = msg.getWindowSize();

            removeLastInWindowS(blkSizeS);

            infoSizeChangedRight();
        }
        else if (msg.getType() == Message.Type.SIZE_CHG) {
            sizeNextWinS = msg.getWindowSize();
        }
        else {
            System.err.println("Illegal message type");
        }
    }

    private void processRight(final Message msg) {
        if (msg.getType() == Message.Type.TUPLE) {
            if (rightMost) {
                if (loadShedder.shedIt(getSizeWindowS())) {
                    logger.info("Tuple S is shed");
                    return;
                }
            }

            Tuple s = msg.getTuple();

            joinWindowR(s);

            addToWindowS(s);

            ackRight();
        }
        else if (msg.getType() == Message.Type.TUPLE_BLK) {
            TupleBlock sBlk = msg.getTupleBlock();

            for (Tuple s : sBlk) {
                joinWindowR(s);

                addToWindowS(s);
            }

            ackBlockRight();
        }
        else if (msg.getType() == Message.Type.ACK) {
            sizeNextWinR = msg.getWindowSize();

            // ignore ack message for R, except for the right-most one
            if (rightMost) {
                evictExpiredTuplesInWindowR();
                infoSizeChangedLeft();
            }
        }
        else if (msg.getType() == Message.Type.ACK_BLK) {
            sizeNextWinR = msg.getWindowSize();

            // ignore ack tuple block message for R
        }
        else if (msg.getType() == Message.Type.SIZE_CHG) {
            sizeNextWinR = msg.getWindowSize();
        }
        else {
            System.err.println("Illegal message type");
        }
    }

    private void forwardTuples() {
        if (winS.size() > sizeNextWinS) {
            Tuple s = (Tuple) winS.remove();
            if (downStreamLeft != null) {
                Message tupleMsg = new Message(getId(),
                        Message.Direction.FROM_RIGHT, s);
                downStreamLeft.put(tupleMsg);
            }

            fwdS.add(s);
        }

        if (winR.size() > sizeNextWinR) {
            Tuple r = (Tuple) winR.remove();
            if (downStreamRight != null) {
                Message tupleMsg = new Message(getId(),
                        Message.Direction.FROM_LEFT, r);
                downStreamRight.put(tupleMsg);
            }

            if (rightMost) {
                fwdR.add(r);
            }
            else {
                if (multimapJoin) {
                    mmapR.remove(r.getTargetField(), r);
                }
                if (bstJoin) {
                    bstR.remove(r.getTargetField(), r);
                }

                infoSizeChangedLeft();
            }
        }
    }

    private void forwardTupleBlocks() {
        if (winS.size() - sizeNextWinS >= blkSizeS) {
            TupleBlock sBlk = new TupleBlock(blkSizeS);

            for (int i = 0; i < blkSizeS; ++i) {
                Tuple s = (Tuple) winS.remove();
                sBlk.add(s);
                fwdS.add(s);
            }

            if (downStreamLeft != null) {
                Message tupleBlkMsg = new Message(getId(),
                        Message.Direction.FROM_RIGHT, sBlk);
                downStreamLeft.put(tupleBlkMsg);
            }
        }

        if (winR.size() - sizeNextWinR >= blkSizeR) {
            TupleBlock rBlk = new TupleBlock(blkSizeR);

            for (int i = 0; i < blkSizeR; ++i) {
                Tuple r = (Tuple) winR.remove();
                rBlk.add(r);
                if (rightMost) {
                    fwdR.add(r);
                }
                else {
                    if (multimapJoin) {
                        mmapR.remove(r.getTargetField(), r);
                    }
                    if (bstJoin) {
                        bstR.remove(r.getTargetField(), r);
                    }
                }
            }

            if (downStreamRight != null) {
                Message tupleBlkMsg = new Message(getId(),
                        Message.Direction.FROM_LEFT, rBlk);
                downStreamRight.put(tupleBlkMsg);
            }

            if (!rightMost) {
                infoSizeChangedLeft();
            }
        }
    }

    private void addToWindowR(final Tuple tuple) {
        winR.add(tuple);

        if (multimapJoin) {
            mmapR.put(tuple.getTargetField(), tuple);
        }
        if (bstJoin) {
            bstR.put(tuple.getTargetField(), tuple);
        }
    }

    private void addToWindowS(final Tuple tuple) {
        winS.add(tuple);

        if (multimapJoin) {
            mmapS.put(tuple.getTargetField(), tuple);
        }
        if (bstJoin) {
            bstS.put(tuple.getTargetField(), tuple);
        }
    }

    private void joinWindowR(final Tuple s) {
        if (multimapJoin) {
            searchMultimapAndJoinWindowR(s);
        }
        else if (bstJoin) {
            searchBSTAndJoinWindowR(s);
        }
        else {
            scanAndJoinWindowR(s);
        }
    }

    private void joinWindowS(final Tuple r) {
        if (multimapJoin) {
            searchMultimapAndJoinWindowS(r);
        }
        else if (bstJoin) {
            searchBSTAndJoinWindowS(r);
        }
        else {
            scanAndJoinWindowS(r);
        }
    }

    private void scanAndJoinWindowR(final Tuple s) {
        for (Object r : winR) {
            if (s.matches((Tuple) r)) {
                emitJoinResult((Tuple) r, s);
            }
        }
        cmpCount += winR.size();

        if (rightMost) {
            for (Tuple r : fwdR) {
                if (s.matches(r)) {
                    emitJoinResult(r, s);
                }
            }
            cmpCount += fwdR.size();
        }
    }

    private void scanAndJoinWindowS(final Tuple r) {
        for (Object s : winS) {
            if (r.matches((Tuple) s)) {
                emitJoinResult(r, (Tuple) s);
            }
        }

        for (Tuple s : fwdS) {
            if (r.matches(s)) {
                emitJoinResult(r, s);
            }
        }

        cmpCount += winS.size() + fwdS.size();
    }

    private void searchMultimapAndJoinWindowR(final Tuple s) {
        Collection<Tuple> matchTuples = mmapR.get(s.getTargetField());

        for (Tuple r : matchTuples) {
            emitJoinResult(r, s);
        }

        cmpCount += 1 + matchTuples.size();
    }

    private void searchMultimapAndJoinWindowS(final Tuple r) {
        Collection<Tuple> matchTuples = mmapS.get(r.getTargetField());

        for (Tuple s : matchTuples) {
            emitJoinResult(r, s);
        }

        cmpCount += 1 + matchTuples.size();
    }

    private void searchBSTAndJoinWindowR(final Tuple s) {
        List<Tuple> matchTuples = bstR.get(s.getTargetField());

        cmpCount += bstR.height();

        if (matchTuples != null) {
            for (Tuple r : matchTuples) {
                emitJoinResult(r, s);
            }

            cmpCount += matchTuples.size();
        }
    }

    private void searchBSTAndJoinWindowS(final Tuple r) {
        List<Tuple> matchTuples = bstS.get(r.getTargetField());

        cmpCount += bstS.height();

        if (matchTuples != null) {
            for (Tuple s : matchTuples) {
                emitJoinResult(r, s);
            }

            cmpCount += matchTuples.size();
        }
    }

    private Tuple peekLastInWindowR() {
        Tuple tuple = null;

        if (rightMost) {
            tuple = fwdR.peek();
            if (tuple == null && !winR.isEmpty()) {
                tuple = (Tuple) winR.get();
            }
        }
        else {
            if (!winR.isEmpty()) {
                tuple = (Tuple) winR.get();
            }
        }

        return tuple;
    }

    private Tuple peekLastInWindowS() {
        Tuple tuple = fwdS.peek();
        if (tuple == null && !winS.isEmpty()) {
            tuple = (Tuple) winS.get();
        }

        return tuple;
    }

    private Tuple removeLastInWindowR() {
        Tuple tuple = null;

        if (rightMost) {
            if (fwdR.isEmpty()) {
                if (!winR.isEmpty()) {
                    tuple = (Tuple) winR.remove();
                }
            }
            else {
                tuple = fwdR.remove();
            }
        }
        else {
            if (!winR.isEmpty()) {
                tuple = (Tuple) winR.remove();
            }
        }

        if (tuple != null) {
            if (multimapJoin) {
                mmapR.remove(tuple.getTargetField(), tuple);
            }
            if (bstJoin) {
                bstR.remove(tuple.getTargetField(), tuple);
            }
        }

        return tuple;
    }

    private void removeLastInWindowR(final int numTuples) {
        for (int i = 0; i < numTuples; ++i) {
            removeLastInWindowR();
        }
    }

    private Tuple removeLastInWindowS() {
        Tuple tuple = null;

        if (fwdS.isEmpty()) {
            if (!winS.isEmpty()) {
                tuple = (Tuple) winS.remove();
            }
        }
        else {
            tuple = fwdS.remove();
        }

        if (tuple != null) {
            if (multimapJoin) {
                mmapS.remove(tuple.getTargetField(), tuple);
            }
            if (bstJoin) {
                bstS.remove(tuple.getTargetField(), tuple);
            }
        }

        return tuple;
    }

    private void removeLastInWindowS(final int numTuples) {
        for (int i = 0; i < numTuples; ++i) {
            removeLastInWindowS();
        }
    }

    private void evictExpiredTuplesInWindowR() {
        Tuple r;
        while ((r = peekLastInWindowR()) != null) {
            if (r.isExpired(currTime)) {
                removeLastInWindowR();
            }
            else {
                break;
            }
        }
    }

    private void evictExpiredTuplesInWindowS() {
        Tuple s;
        while ((s = peekLastInWindowS()) != null) {
            if (s.isExpired(currTime)) {
                removeLastInWindowS();
            }
            else {
                break;
            }
        }
    }

    private void ackLeft() {
        if (downStreamLeft != null) {
            int windowSize = getSizeWindowR();
            Message ackMsg = new Message(getId(), Message.Type.ACK,
                    Message.Direction.FROM_RIGHT, windowSize);

            downStreamLeft.put(ackMsg);
        }
    }

    private void ackRight() {
        if (downStreamRight != null) {
            int windowSize = getSizeWindowS();
            Message ackMsg = new Message(getId(), Message.Type.ACK,
                    Message.Direction.FROM_LEFT, windowSize);

            downStreamRight.put(ackMsg);
        }
    }

    private void ackBlockLeft() {
        if (downStreamLeft != null) {
            int windowSize = getSizeWindowR();
            Message ackMsg = new Message(getId(), Message.Type.ACK_BLK,
                    Message.Direction.FROM_RIGHT, windowSize);

            downStreamLeft.put(ackMsg);
        }
    }

    private void ackBlockRight() {
        if (downStreamRight != null) {
            int windowSize = getSizeWindowS();
            Message ackMsg = new Message(getId(), Message.Type.ACK_BLK,
                    Message.Direction.FROM_LEFT, windowSize);

            downStreamRight.put(ackMsg);
        }
    }

    private void infoSizeChangedLeft() {
        if (downStreamLeft != null) {
            int windowSize = getSizeWindowR();
            Message ackMsg = new Message(getId(), Message.Type.SIZE_CHG,
                    Message.Direction.FROM_RIGHT, windowSize);

            downStreamLeft.put(ackMsg);
        }
    }

    private void infoSizeChangedRight() {
        if (downStreamRight != null) {
            int windowSize = getSizeWindowS();
            Message ackMsg = new Message(getId(), Message.Type.SIZE_CHG,
                    Message.Direction.FROM_LEFT, windowSize);

            downStreamRight.put(ackMsg);
        }
    }

    private int getSizeWindowS() {
        return winS.size() + fwdS.size();
    }

    private int getSizeWindowR() {
        int windowSizeR = winR.size();
        if (rightMost) {
            windowSizeR += fwdR.size();
        }

        return windowSizeR;
    }

    private void emitJoinResult(final Tuple r, final Tuple s) {
        JoinResult result = new JoinResult(getId(), r, s);

        downStreamResult.put(result);
    }

    public JoinWorker setDownStreamRight(Stream<Message> stream) {
        this.downStreamRight = Preconditions.checkNotNull(stream);

        return this;
    }

    public JoinWorker setDownStreamLeft(Stream<Message> stream) {
        this.downStreamLeft = Preconditions.checkNotNull(stream);

        return this;
    }

    public JoinWorker setDownStreamResult(Stream<JoinResult> stream) {
        this.downStreamResult = Preconditions.checkNotNull(stream);

        return this;
    }

    public JoinWorker setRightMost(final boolean rightMost) {
        this.rightMost = rightMost;

        return this;
    }

    public JoinWorker setLeftMost(final boolean leftMost) {
        this.leftMost = leftMost;

        return this;
    }

    public JoinWorker setCapacityOfWindowR(final int size) {
        this.capWinR = size;

        return this;
    }

    public JoinWorker setCapacityOfWindowS(final int size) {
        this.capWinS = size;

        return this;
    }

    public JoinWorker setEnableBlockedTransfer(final boolean blockedTrans) {
        this.blockedTrans = blockedTrans;

        return this;
    }

    public JoinWorker setTupleBlockSizeR(final int size) {
        this.blkSizeR = size;

        return this;
    }

    public JoinWorker setTupleBlockSizeS(final int size) {
        this.blkSizeS = size;

        return this;
    }

    public JoinWorker setEnableMultimapJoin(final boolean multimapJoin) {
        this.multimapJoin = multimapJoin;

        return this;
    }

    public JoinWorker setEnableBSTJoin(final boolean bstJoin) {
        this.bstJoin = bstJoin;

        return this;
    }

    public JoinWorker setWorkerName(final String name) {
        setName(name);

        return this;
    }

    public JoinWorker setDisplayStatusInterval(final long interval) {
        if (interval >= 0) {
            setTimerInterval(interval, TimeUnit.MILLISECONDS);
        }

        return this;
    }

    private String getDisplayName() {
        return getName() + "-" + getId();
    }

    private String getStat() {
        StringBuilder sb = new StringBuilder();

        sb.append(getDisplayName() + ": ");

        Joiner joinerComma = Joiner.on(", ").skipNulls();
        List<String> stat = Lists.newLinkedList();

        stat.add("winR=" + winR.size());
        if (rightMost) {
            stat.add("fwdR=" + fwdR.size());
        }
        stat.add("winS=" + winS.size());
        stat.add("fwdS=" + fwdS.size());
        if (multimapJoin) {
            stat.add("mmapR=" + mmapR.size());
            stat.add("mmapS=" + mmapS.size());
        }
        if (bstJoin) {
            stat.add("bstR(keys)=" + bstR.size());
            stat.add("bstS(keys)=" + bstS.size());
        }

        stat.add("cmpCount=" + cmpCount);

        sb.append(joinerComma.join(stat));

        return sb.toString();
    }
}